from typing import Optional, Tuple, List
from typeguard import check_argument_types
from pathlib import Path
import numpy as np
import librosa
import onnxruntime # 1.12.1
import logging
import os


def mask_fill(arr, mask, mask_value):
    """Numpy implementation of torch.Tensor.masked_fill

    Args:
        arr (np.ndarray): Array to mask.
        mask (np.ndarray): Every value with 1 (or True) will mask with mask_value
        mask_value (number): Value to apply for mask

    Returns:
        np.ndarray: Masked array
    """
    arr[mask.astype(np.bool) == True] = mask_value
    return arr


def make_pad_mask(lengths, xs=None, dim=-1, xs_shape=None):
    """Make mask tensor containing indices of padded part.
    
    Args:
        lengths (np.ndarray or List): Batch of lengths (B,).
        xs (np.ndarray, optional): The reference tensor.
            If set, masks will be the same shape as this tensor.
        dim (int, optional): Dimension indicator of the above tensor.
            See the example.

    Returns:
        np.ndarray: Mask tensor containing indices of padded part.

    Examples:
        With only lengths.
        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        masks = [[0, 0, 0, 0 ,0],
                 [0, 0, 0, 1, 1],
                 [0, 0, 1, 1, 1]]
                 
        With the reference tensor.
        >>> xs = np.zeros((3, 2, 4))
        >>> make_pad_mask(lengths, xs)
        array([[[0, 0, 0, 0],
                 [0, 0, 0, 0]],
                [[0, 0, 0, 1],
                 [0, 0, 0, 1]],
                [[0, 0, 1, 1],
                 [0, 0, 1, 1]]])
                 
        >>> xs = np.zeros((3, 2, 6))
        >>> make_pad_mask(lengths, xs)
        array([[[0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1]],
                [[0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1]],
                [[0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1]]])
                 
        With the reference tensor and dimension indicator.
        >>> xs = np.zeros((3, 6, 6))
        >>> make_pad_mask(lengths, xs, 1)
        array([[[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1]],
                [[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1]],
                [[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1]]])
                 
        >>> make_pad_mask(lengths, xs, 2)
        array([[[0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1]],
                [[0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1]],
                [[0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1]]])
    """
    if xs is not None:
        base = np.zeros(xs.shape)
    else:
        base = np.zeros((len(lengths), max(lengths)))

    if len(base.shape) == 3 and dim == 1:
        base = base.transpose(0, 2, 1)

    for i in range(len(base)):
        base[i][..., lengths[i]:] = 1

    if len(base.shape) == 3 and dim == 1:
        base = base.transpose(0, 2, 1)

    return base


"""
     audio
       🠇
       🠇
extract features
       🠇
       🠇
  normalization
       🠇
       🠇
    encoder
"""

def frontend_stft(input: np.ndarray, ilens: np.ndarray = None) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    """STFT forward function.
    Args:
        input: (Batch, Nsamples)
        ilens: (Batch)
    Returns:
        output: (Batch, Frames, Freq, 2)
    """
    assert check_argument_types()
    n_fft, hop_length = 512, 128
    stft_kwargs = dict(
        n_fft=512, #  config.n_fft,
        win_length=512, # config.win_length,
        hop_length=128, #s config.hop_length,
        center=True, # config.center,
        window="hann" # config.window,
    )
    output = []
    # iterate over istances in a batch
    for i, instance in enumerate(input):
        stft = librosa.stft(input[i], **stft_kwargs)
        output.append(np.array(np.stack([stft.real, stft.imag], -1)))
    output = np.vstack(output).reshape(len(output), *output[0].shape)

    onesided, normalized, center  = True, False, True

    if not onesided: #self.config.onesided:
        len_conj = n_fft - output.shape[1]
        conj = output[:, 1: 1 + len_conj].flip(1)
        conj[:, :, :, -1].data *= -1
        output = np.concatenate([output, conj], 1)
    if normalized: # self.config.normalized:
        output = output * (stft_kwargs["window"].shape[0] ** (-0.5))

    # output: (Batch, Freq, Frames, 2=real_imag)
    # -> (Batch, Frames, Freq, 2=real_imag)
    output = output.transpose(0, 2, 1, 3)
    if ilens is not None:
        if center: # self.config.center:
            pad = n_fft // 2
            ilens = ilens + 2 * pad
        olens = (ilens - n_fft) // hop_length + 1
        output = mask_fill(output, make_pad_mask(
            olens, output, dim=1), 0.0)
    else:
        olens = None
    return output, olens


def frontend_logmel(feat: np.ndarray, ilens: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray]:
    mel_options = dict(sr=16000, n_fft=512, n_mels=80, fmin=0, fmax=8000.0, htk=False)
    log_base = None
    melmat = librosa.filters.mel(**mel_options)
    # self.melmat = melmat.T

    # feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
    mel_feat = np.matmul(feat, melmat.T)#self.melmat)
    mel_feat = np.clip(mel_feat, 1e-10, float('inf'))

    if log_base is None:
        logmel_feat = np.log(mel_feat)
    elif log_base == 2.0:
        logmel_feat = np.log2(mel_feat)
    elif log_base == 10.0:
        logmel_feat = np.log10(mel_feat)
    else:
        logmel_feat = np.log(mel_feat) / np.log(log_base)

    # Zero padding
    logmel_feat = mask_fill(logmel_feat, make_pad_mask(ilens, logmel_feat, 1), 0.0)
    return logmel_feat, ilens


def extract_features(inputs: np.ndarray, input_length: np.ndarray):
    assert check_argument_types()
    # STFT
    input_stft, feats_lens = frontend_stft(inputs, input_length)
    # 3. STFT -> Power spectrum
    # h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
    input_power = input_stft[..., 0]**2 + input_stft[..., 1]**2

    # 4. Feature transform e.g. Stft -> Log-Mel-Fbank
    # input_power: (Batch, [Channel,] Length, Freq)
    #       -> input_feats: (Batch, Length, Dim)
    input_feats, _ = frontend_logmel(input_power, feats_lens)
    return input_feats, feats_lens

## End for Feature Extract

def global_mvn(x: np.ndarray, ilens: np.ndarray, model_dir: str) -> Tuple[np.ndarray, np.ndarray]:

    stats_file = "/home/gyf/.cache/espnet_onnx/conformer/feats_stats.npz"
    stats_file = os.path.join(model_dir, "feats_stats.npz")
    stats_file = Path(stats_file)
    stats = np.load(stats_file)
    eps = 1.0e-20
    if isinstance(stats, np.ndarray):
        # Kaldi like stats
        count = stats[0].flatten()[-1]
        mean = stats[0, :-1] / count
        var = stats[1, :-1] / count - mean * mean
    else:
        # New style: Npz file
        count = stats["count"]
        sum_v = stats["sum"]
        sum_square_v = stats["sum_square"]
        mean = sum_v / count
        var = sum_square_v / count - mean * mean
    std = np.sqrt(np.maximum(var, eps))

    """Forward function
    Args:
        x: (B, L, ...)
        ilens: (B,)
    """
    mask = make_pad_mask(ilens, x, 1)

    # feat: (B, T, D)
    if True: #norm_means:
        x -= mean
    x = mask_fill(x, mask, 0.0)

    if True: #norm_vars:
        x /= std

    return x, ilens


def read_audio(wav_file):
    speech, sr = librosa.load(str(wav_file), sr=16000)
    # check dtype
    if speech.dtype != np.float32:
        speech = speech.astype(np.float32)

    # data: (Nsamples,) -> (1, Nsamples)
    speech = speech[np.newaxis, :]
    # lengths: (1,)
    lengths = np.array([speech.shape[1]]).astype(np.int64)
    return speech, lengths


def frontend_main(wav_file: str, model_dir: str) -> Tuple[np.ndarray, np.ndarray]:
    # speech: np.ndarray, speech_length: np.ndarray
    speech, speech_length = read_audio(wav_file)
    feats, feat_length = extract_features(speech, speech_length)
    feats, feat_length = global_mvn(feats, feat_length, model_dir)
    return feats, feat_length

## End for Frontend_main
